import cv2 as cv
import numpy as np
import apriltag
from math import *
from uarm.wrapper import SwiftAPI
import matplotlib.pyplot as plt


in_pose = []
in_angle = []

# 相机内参矩阵
mtx = np.load('./data/camera_param.npz')['mtx']
dist = np.load('./data/camera_param.npz')['dist']

base2tag_p = np.array([[0.175, 0.005, -0.07],
        [0.174, -0.0163, -0.0727],
        [0.1447, 0.0056, -0.072],
        [0.1438, -0.0144, -0.0728]])
base2tag_p = np.array([[175, 5, -70],
        [174, -16.3, -72.7],
        [144.7, 5.6, -72],
        [143.8, 14.4, -72.8]])
# 定义一个检测器（使用字典“tag36h10”）
at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11 tag36h10'))


def eulerAngles2RotationMatrix(theta):
    """
    欧拉角转换旋转矩阵
    :param theta: 欧拉角角度
    :return: 旋转矩阵
    """
    R_x = np.array([[1, 0, 0],
                    [0, cos(theta[0]), sin(theta[0])],
                    [0, -sin(theta[0]), cos(theta[0])]
                    ])
    R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],
                    [0, 1, 0],
                    [-sin(theta[1]), 0, cos(theta[1])]
                    ])
    R_z = np.array([[cos(theta[2]), -sin(theta[2]),0],
                    [sin(theta[2]), cos(theta[2]),0],
                    [0, 0, 1]
                    ])
    R = R_z@R_y@R_x
    return R


def rotationMatrixToEulerAngles(R):
    """
    旋转矩阵转换欧拉角
    :param R: 旋转矩阵
    :return: 欧拉角
    """
    sy = sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = atan2(R[2, 1], R[2, 2])
        y = atan2(-R[2, 0], sy)
        z = atan2(R[1, 0], R[0, 0])
    else:
        x = atan2(-R[1, 2], R[1, 1])
        y = atan2(-R[2, 0], sy)
        z = 0
    return np.array([x, y, z])

def get_tag2end(point,position):
    """
    获取二维码到手臂末端
    :param point:
    :return:
    """

    # 计算绕z轴旋转角
    obj= point[:,:2]
    delta = (obj[2]-obj[3])
    r = sqrt(delta@delta.T)
    angle = np.arcsin(delta[0]/r)
    # 计算变换后的坐标
    R=eulerAngles2RotationMatrix([0,0,angle])
    t = R @ np.array(position) - R @ point[3]
    return R,t


def get_img2tag(image,tag_size):
    img2tag_p = []
    cam_params = [mtx[0, 0], mtx[1, 1], mtx[0, 2], mtx[1, 2]]
    gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
    tags = at_detector.detect(gray)
    for tag in tags:
        M, e1, e2 = at_detector.detection_pose(tag, cam_params, tag_size)
        img2tag_p.append([tag.tag_id,tag.center,M])
    return img2tag_p


def get_tag2cam(tags):
    ret = False
    t = []
    R = []
    obj_point = np.array([[30, 30, 0],
                        [30, 0, 0],
                        [0, 30, 0],
                        [0, 0, 0]], dtype=float)
    obj_id = [0,1,2,3]
    pix_point = np.zeros((4, 2))
    for tag in tags:
        for id in obj_id:
            if tag[0] == id:
                pix_point[id, :] = tag[1]
                obj_id.remove(id)
    if len(obj_id)==0:
        pix_point = np.array(pix_point)
        obj_point = np.array(obj_point)
        ret, rvecs, tvecs = cv.solvePnP(obj_point, pix_point, mtx, dist)
        R = cv.Rodrigues(rvecs)[0]
        t = tvecs.T
        ret = True
    return ret,R,t

def get_base2end(angle,position):
    R = eulerAngles2RotationMatrix([0,0,(angle)*pi/180])
    t = position
    return R,t

def get_base2tag_p(tags,id,R,t,angle,position):
    ret = False
    dst = []
    for tag in tags:
        if tag[0]==id:
            pix2img_t = np.linalg.inv(mtx) @ np.insert(tag[1], 2, 1)
            img2tag_t = (tag[2] @ np.insert(pix2img_t, 3, 1))[:3]
            base2end_R, base2end_t = get_base2end(angle,position)
            dst = base2end_t + base2end_R @ R @ img2tag_t + base2end_R @ -R@t
            ret = True
            print(dst)
    return ret,dst

if __name__ == "__main__":
    frame = cv.imread("./test.jpeg")

    position=[139.8,1.8,84.7]
    angle = 0.7676
    cv.imshow("frame", frame)

    tags = get_img2tag(frame, 20)
    ret, tag2cam_R, tag2cam_t = get_tag2cam(tags)
    tag2end_R, tag2end_t = get_tag2end(base2tag_p, position)
    if ret is True:
        base2end_R, base2end_t = get_base2end(0,position)
        cam2end_R = tag2cam_R.T @ tag2end_R @ base2end_R
        cam2end_t = (tag2cam_R @ tag2end_t + tag2cam_t)[0]
        np.savez("./data/eyehand_Matrix.npz", R=cam2end_R, t=cam2end_t)
        print(f"save as ./data/eyehand_Matrix.npz")
        print(f"旋转矩阵转换欧拉角{rotationMatrixToEulerAngles(cam2end_R)}")
        print(f"平移矩阵{cam2end_t}")
    else:
        print("未检测到二维码")

